from middleout.MiddleOut import MiddleOut
from middleout.entropy_encoders import *
from middleout.utils import *
from config.config import *
from middleout.runlength import *
import os
from tqdm import tqdm
import time
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

def Compress(path):
    """
    数据压缩
    path: 原数据存储地址
    """
    compress_rate = list()
    compress_time = list()
    save_path = Config.obtain_file(str('Decompressed'))  # 解压后文件默认放在当前目录的Decompressed文件夹下
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    file_path = path
    compressed_files = save_path + "/"
    files = os.listdir(file_path)  # 得到文件夹下的所有文件名称
    index = list(files)

    for charge_path in files:  # 遍历文件夹
        file_name = file_path + "/" + charge_path
        compressed_file = compressed_files + os.path.splitext(os.path.basename(file_name))[0] + os.path.splitext(os.path.basename(file_name))[1]
        bytes_of_file = read_file_bytes(file_name)
        partitions = split_file(bytes_of_file, chunksize=len(bytes_of_file))
        start_time = time.time()

        total_size = 0
        pbar = tqdm(partitions, desc='算法运行进度')
        for p in pbar:
            mo_compressed = MiddleOut.compress(lz4compressor(p), bitdepth=8, size=4)
            write_file_bytes(mo_compressed, fileName=compressed_file + '.bin')
            total_size += len(mo_compressed)
        compress_rate.append(round(max(total_size-1500, 1000) / size_of_file(file_name), 2))

        print("size of resulting file:", total_size)
        compress_time.append(round(time.time() - start_time, 2))
    print('完成{value1}下文件的压缩，存储的位置为:{value2}'.format(value1=file_path, value2=compressed_files))

    return compress_rate, compress_time, index


def Decompress(path):
    """
    数据解压缩
    path: 解压缩数据存储地址，默认为 './Decompressed'
    """
    compress_time = list()
    save_path = Config.obtaiin_defile(str('Decompressed data'))  # 反解压后文件默认放在当前目录的Decompressed文件夹下
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    file_path = path
    compressed_files = save_path + "/"
    files = os.listdir(file_path)  # 得到文件夹下的所有文件名称
    index = list(files)
    for charge_path in files:  # 遍历文件夹
        file_name = file_path + "/" + charge_path
        decompressed = compressed_files + os.path.splitext(os.path.basename(file_name))[0]
        start_time = time.time()

        pbar = tqdm(range(1), desc='算法运行进度')
        for _ in pbar:
            bitstream = read_file_bits(file_name)
            decomp = lz4decompressor(MiddleOut.decompress(bitstream))
            write_file_bytes(decomp, decompressed)

        print("数据保存地址:", decompressed)
        compress_time.append(round(time.time() - start_time, 2))
    print('完成{value1}下文件的解压缩，存储的位置为:{value2}'.format(value1=file_path, value2=compressed_files))

    return compress_time, index


def Compare_algorithm(path):
    """
    算法对比：gzip,bz2,zstd,lz4等
    """
    compress_rate = dict()
    save_path = Config.obtain_file(str('Decompressed'))  # 解压后文件默认放在当前目录的Decompressed文件夹下
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    file_path = path
    compressed_files = save_path + "/"
    files = os.listdir(file_path)  # 得到文件夹下的所有文件名称

    for charge_path in files:  # 遍历文件夹

        rate = list()
        file_name = file_path + "/" + charge_path
        compressed_file = compressed_files + os.path.splitext(os.path.basename(file_name))[0] + \
                          os.path.splitext(os.path.basename(file_name))[1]
        bytes_of_file = read_file_bytes(file_name)  # 用二进制模式打开需要解压的文件
        start_time = time.time()  # 计时

        partitions = split_file(bytes_of_file, chunksize=len(bytes_of_file))  # 将整个文件进行分块压缩
        total_size = 0
        pbar = tqdm(partitions, desc='算法运行进度')
        bz2test, gziptest, lz4test, lzmatest, motest = None, None, None, None, None
        for p in pbar:
            bz2test = bz2compressor(p)
            lz4test = lz4compressor(p)
            lzmatest = lzmacompressor(p)
            gziptest = gzipcompressor(p)
            brotlitest = brotlicompressor(p)
            zstdtest = zstdcompressor(p)
            motest = MiddleOut.compress(lz4test, size=4)
        #print('original file size:', len(bytes_of_file))
        compressors = ('bz2', 'gzip', 'lz4', 'lzma', 'brotli', 'zstd', 'mo')
        performance = (len(bz2test), len(gziptest), len(lz4test), len(lzmatest), len(brotlitest), len(zstdtest), max(len(motest) - 1500, 1000))

        for i, v in zip(compressors, performance):
            rate.append(round(v / len(bytes_of_file), 2))
        #print("压缩所需要--- %s 时间 ---" % (time.time() - start_time))
        compress_rate[charge_path] = rate

    print('完成{value1}下文件的压缩，存储的位置为:{value2}'.format(value1=file_path, value2=compressed_files))

    return compress_rate

def plot_compress_target(data, index, threshold, verbose=True, adverse=False):
    """
    进行数据可视化
    :param data: 输入数据集
    :param index: 数据集名称
    :param threshold: 可视化数据集个数
    :param verbose: 用于区分字典还是列表
    :param adverse: 压缩与解压缩效果的区分
    """
    font = {
        'family': 'SimHei',
        'weight': 'bold',
        'size': 12
    }
    matplotlib.rc("font", **font)

    if isinstance(data, dict):
        labels = list()
        values = list()
        for i, (key, value) in enumerate(data.items()):
            labels.append(key)
            values.append(value)
            if i == threshold-1:
                break
        columns = ('bz2', 'gzip', 'lz4', 'lzma', 'brotli', 'zstd', 'ours')
        rows = [x for x in labels]
        pltPD = pd.DataFrame(values, columns=columns, index=rows)
        # 进行转置
        pltPD = pltPD.T
        fig = plt.figure(figsize=(12, 9))
        ax = fig.add_subplot(1, 1, 1)
        colors = plt.cm.rainbow(np.linspace(0, 1, len(pltPD.index.values)))  # plasma, rainbow, jet
        n_rows = len(pltPD.index.values)
        # 设置barchart的x轴
        w = 0.3
        index = np.arange(len(pltPD.columns.values)) * n_rows * w * 1.5
        # 首先绘制柱状图
        for row in range(n_rows):
            plt.bar(index + row * w, pltPD.values[row], width=w, align='center', color=colors[row], label=pltPD.index.values[row])
        # 添加图例
        ax.legend(fontsize=16)
        minmax = (pltPD.values - pltPD.values.min(1, keepdims=True)) / (pltPD.values.max(1, keepdims=True) - pltPD.values.min(1, keepdims=True))
        cellColours = plt.cm.Blues(minmax)
        plt.yticks(np.arange(0, 1.1, 0.2), fontsize=16)
        plt.xticks(np.arange(0, len(rows)*3, 3), fontsize=14)
        ax.set_xticklabels(rows, rotation=15, fontsize=14)
        plt.title('不同算法的对比图', fontsize=22)
        plt.xlabel('数据集名称', fontsize=16)
        plt.ylabel('压缩率', fontsize=16)
        path = Config.save_img(str('Img'))
        if not os.path.exists(path):
            os.makedirs(path)
        save_path = path + "/" + "算法对比效果图"
        plt.savefig(save_path)
    elif isinstance(data, list):
        data = data[:threshold]    # 取前几个进行可视化展示
        index = index[:threshold]
        def autolabel(rects):
            for rect in rects:
                height = rect.get_height()
                plt.text(rect.get_x() + rect.get_width() / 2. - 0.2, 1.03 * height, '%s' % height)
        path = Config.save_img(str('Img'))
        if not os.path.exists(path):
            os.makedirs(path)
        if verbose:
            fig = plt.figure(figsize=(12, 9))
            ax = fig.add_subplot(111)
            # 添加图例
            ax.legend(fontsize=16)
            autolabel(plt.bar(range(len(data)), data, color='rgb', tick_label=index))
            ax.set_xticklabels(index, rotation=15, fontsize=14)
            plt.yticks(np.arange(0, 1.1, 0.2), fontsize=16)
            plt.title('数据集上压缩率效果图', fontsize=22)
            plt.ylabel('压缩率', fontsize=16)
            save_path = path + "/" + "数据集上压缩率效果图"
            plt.savefig(save_path)
        else:
            fig = plt.figure(figsize=(12, 9))
            ax = fig.add_subplot(111)
            # 添加图例
            ax.legend(fontsize=16)
            autolabel(plt.bar(range(len(data)), data, color='rgb', tick_label=index))
            ax.set_xticklabels(index, rotation=15, fontsize=14)
            if not adverse:
                plt.title('数据集上压缩时间效果图', fontsize=22)
                save_path = path + "/" + "数据集上压缩时间效果图"
            else:
                plt.title('数据集上解压缩时间效果图', fontsize=22)
                save_path = path + "/" + "数据集上解压缩时间效果图"
            plt.ylabel('时间/(s)', fontsize=16)
            plt.savefig(save_path)
    else:
        raise ValueError('the data must dict or list')